package com.zzh.zzhaiagent.agent;

import cn.hutool.core.util.StrUtil;
import com.zzh.zzhaiagent.agent.model.AgentState;
import lombok.Data;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;

import java.io.IOException;
import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.CompletableFuture;

/**
 * @author: zzh
 * @date: 2025/05/22 20:31:26
 * @version: 1.0
 */
//Agent基类 (状态切换 + 基本属性 + AgentLoop循环执行步骤，具体执行抽象化)
@Data
@Slf4j
public abstract class BaseAgent {//因为抽象方法所以要变成抽象类
    //核心属性
    private String name;//Agent名称

    //提示词
    private String systemPrompt;//系统提示词
    private String nextPrompt;//下一步执行的提示词

    //Agent状态
    private AgentState state = AgentState.IDLE;//Agent状态 默认空闲

    //执行控制步骤
    private int currentStep = 0;//当前执行步骤
    private int maxStep = 12;//最大执行步骤 (大于12时，Agent跳出执行循环)

    //提供大模型
    private ChatClient chatClient;

    //提供Memory (需要自主维护会话上下文)
    private List<Message> messageList = new ArrayList<>();

    /**
     * Agent执行步骤 (具体步骤内部执行流程由ReActAgent实现)
     * @param userPrompt
     * @return
     */
    public String run(String userPrompt){
        //判断状态
        if (this.state != AgentState.IDLE){
            //todo 后续完成异常统一处理
            throw new RuntimeException("Cannot run agent from state: " + this.state);
        }

        //判断参数
        if(StrUtil.isBlank(userPrompt)){
            throw new RuntimeException("Cannot run agent with empty user prompt");
        }

        //更改状态（开始执行）
        state = AgentState.RUNNING;
        //记录消息上下文
        messageList.add(new UserMessage(userPrompt));

        //保存结果列表
        List<String> results = new ArrayList<>();
        try {
            for (int i = 0; i < maxStep && state != AgentState.FINISHED; i++) {
                int stepNumber = i + 1;
                currentStep = stepNumber;
                log.info("Executing step " + stepNumber + "/" + maxStep);
                // 单步执行
                String stepResult = step();
                String result = "Step " + stepNumber + ": " + stepResult;
                results.add(result);
            }
            //检查是否超出步骤限制
            if (currentStep >= maxStep) {
                state = AgentState.FINISHED;
                results.add("Terminated: Reached max steps (" + maxStep + ")");
            }
            return String.join("\n", results);
        }catch (Exception e){
            state = AgentState.ERROR;
            log.error("Error executing agent", e);
            return "执行错误" + e.getMessage();
        }finally {
            //清理资源
            this.cleanup();
        }
    }


    /**
     * 运行代理（流式输出）
     * 对比run方法，run方法是同步执行，runStream方法是异步执行，返回一个SseEmitter对象，通过SSE方式实时输出执行结果。
     * @param userPrompt 用户提示词
     * @return 执行结果
     */
    public SseEmitter runStream(String userPrompt) {
        // 创建一个超时时间较长的 SseEmitter
        SseEmitter sseEmitter = new SseEmitter(300000L); // 5 分钟超时
        // 使用线程异步处理，避免阻塞主线程 ，要不断返回sseEmitter.send()
        CompletableFuture.runAsync(() -> {
            // 1、基础校验
            try {
                if (this.state != AgentState.IDLE) {
                    sseEmitter.send("错误：无法从状态运行代理：" + this.state);
                    sseEmitter.complete();//完成结束sseEmitter
                    return;
                }
                if (StrUtil.isBlank(userPrompt)) {
                    sseEmitter.send("错误：不能使用空提示词运行代理");
                    sseEmitter.complete();//完成结束sseEmitter
                    return;
                }
            } catch (Exception e) {
                sseEmitter.completeWithError(e);
            }
            // 2、执行，更改状态
            this.state = AgentState.RUNNING;
            // 记录消息上下文
            messageList.add(new UserMessage(userPrompt));
            // 保存结果列表
            List<String> results = new ArrayList<>();
            try {
                // 执行循环
                for (int i = 0; i < maxStep && state != AgentState.FINISHED; i++) {
                    int stepNumber = i + 1;
                    currentStep = stepNumber;
                    log.info("Executing step {}/{}", stepNumber, maxStep);
                    // 单步执行
                    String stepResult = step();
                    String result = "Step " + stepNumber + ": " + stepResult + "\n";
                    results.add(result);
                    // 输出当前每一步的结果到 SSE
                    sseEmitter.send(result);
                }
                // 检查是否超出步骤限制
                if (currentStep >= maxStep) {
                    state = AgentState.FINISHED;
                    results.add("Terminated: Reached max steps (" + maxStep + ")");
                    sseEmitter.send("执行结束：达到最大步骤（" + maxStep + "）");
                }
                // 正常完成
                sseEmitter.complete();
            } catch (Exception e) {
                state = AgentState.ERROR;
                log.error("error executing agent", e);
                try {
                    sseEmitter.send("执行错误：" + e.getMessage());
                    sseEmitter.complete();
                } catch (IOException ex) {
                    sseEmitter.completeWithError(ex);
                }
            } finally {
                // 3、清理资源
                this.cleanup();
            }
        });

        // 设置超时回调
        sseEmitter.onTimeout(() -> {
            this.state = AgentState.ERROR;
            this.cleanup();
            log.warn("SSE connection timeout");
        });
        // 设置完成回调
        sseEmitter.onCompletion(() -> {
            if (this.state == AgentState.RUNNING) {
                this.state = AgentState.FINISHED;
            }
            this.cleanup();
            log.info("SSE connection completed");
        });
        return sseEmitter;
    }


    /**
     * 执行单个步骤 （think拿到MessageList，根据上下文执行）
     * @return
     */
    public abstract String step();


    /**
     * 清理资源
     */
    public void cleanup(){
        //todo 后续完成资源释放
    }


}

